package cn.ai.shoes.client.utils;

import cn.ai.shoes.common.exception.ServiceException;
import cn.ai.shoes.common.exception.constant.GlobalErrorCodeConstants;
import com.alibaba.fastjson2.JSON;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

@Slf4j
public class LoadBalanceUtils {
    // 添加静态计数器用于轮询
    private static final AtomicInteger currentIndex = new AtomicInteger(0);
    public static final int TIMEOUT = 2;

    public static String switchByGPU(Map<String, String> serverUrls) {
        Collection<String> urlList = serverUrls.values();
        if (urlList.isEmpty()) {
            throw new ServiceException(GlobalErrorCodeConstants.SERVER_NOT_EXIST);
        }
        Map<String,Integer> treeMap =new TreeMap<>();
        for (String server : urlList) {
            if (!server.startsWith("http")) {
                server = "https://" + server;
            }
            if (server.endsWith("/")) {
                server = server.substring(0, server.length() - 1);
            }
            OkHttpClient client = new OkHttpClient().newBuilder()
                    .connectTimeout(TIMEOUT, TimeUnit.SECONDS)
                    .readTimeout(TIMEOUT, TimeUnit.SECONDS)
                    .build();
            Request request = new Request.Builder()
                    .url(server + "/queue")
                    .get()
                    .build();
            try  (Response response = client.newCall(request).execute()){
                if (!response.isSuccessful()) {
                    log.info("请求失败，状态码: " + response.code()+"，响应体: " + response.body().string());
                    continue;
                }else {
                    try (ResponseBody body = response.body()) {
                        if (body != null) {
                            String jsonStr = body.string();
                            Map<String, List<String>> jsonMap = JSON.parseObject(jsonStr, Map.class);
                            List<String> queueRunning =jsonMap.get("queue_running");
                            int runningCount = queueRunning != null ? queueRunning.size() : 0;
                            List<String> queuePending =  jsonMap.get("queue_pending");
                            int pendingCount = queuePending != null ? queuePending.size() : 0;
                            int totalCount = runningCount + pendingCount;
                            treeMap.put(server,totalCount);
                        }else {
                            log.info("响应体为空");
                        }
                    }
                }
            } catch (IOException e) {
                System.err.println("请求异常: " + e.getMessage());
                e.printStackTrace();
                continue;
            }
        }
        if (treeMap.isEmpty()) {
            throw new ServiceException(GlobalErrorCodeConstants.SERVER_NOT_EXIST);
        }
        List<Map.Entry<String, Integer>> collect = treeMap.entrySet().stream()
                .sorted(Map.Entry.comparingByValue())
                .collect(Collectors.toList());
        String server = collect.get(0).getKey();
        return server;
    }
    /**
     * 轮询
     */
    public static String switchByPolling(Map<String, String> serverUrls) {
        Collection<String> urlList = serverUrls.values();
        List<String> aliveServers = new ArrayList<>();
        for (String url : urlList) {
            if (isServerAlive(url)) {
                aliveServers.add(url);
            }
        }
        if (aliveServers.isEmpty()) {
            throw new ServiceException(GlobalErrorCodeConstants.SERVER_NOT_EXIST);
        }
        int index = currentIndex.getAndIncrement() % aliveServers.size();
        if (index < 0) {
            index += aliveServers.size();
            currentIndex.set(0);
        }
        String server = aliveServers.get(index);
        if (!server.startsWith("http")) {
            server = "https://" + server;
        }
        return server;
    }

    public static String notSwitch(Map<String, String> serverUrls, Integer number) {
        Collection<String> urlList = serverUrls.values();
        List<String> aliveServers = new ArrayList<>();
        for (String url : urlList) {
            if (isServerAlive(url)) {
                aliveServers.add(url);
            }
        }
        if (aliveServers.isEmpty()) {
            throw new ServiceException(GlobalErrorCodeConstants.SERVER_NOT_EXIST);
        }
        if (number >= aliveServers.size()) {
            throw new ServiceException(GlobalErrorCodeConstants.SERVER_NOT_EXIST);
        }
        return aliveServers.get(number);
    }
    /**
     * 检查服务器心跳
     */
    public static boolean isServerAlive(String serverUrl) {
        // 去掉结尾的斜杠，避免出现 // 的情况
        if (serverUrl.endsWith("/")) {
            serverUrl = serverUrl.substring(0, serverUrl.length() - 1);
        }
        if (!serverUrl.startsWith("http")) {
            serverUrl = "https://" + serverUrl;
        }
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(TIMEOUT, TimeUnit.SECONDS)
                .readTimeout(TIMEOUT, TimeUnit.SECONDS)
                .build();
        Request request = new Request.Builder()
                .url(serverUrl)
                .get()
                .build();
        try (Response response = client.newCall(request).execute()) {
            return response.isSuccessful();
        } catch (IOException e) {
            System.out.println("心跳检测失败: " + serverUrl + "，异常信息: " + e.getMessage());
            return false;
        }
    }

}
